use std::{pin::Pin, sync::{Arc, Mutex}, io::Write};
use futures::Stream;
use tokio::{
    io::{AsyncBufReadExt, BufReader as TokioBufReader},
    sync::broadcast,
};
use async_trait::async_trait;

use crate::{McpError, McpErrorCode, McpResult};
use super::{McpMessage, McpTransport};

pub struct McpStdioTransport {
    /// 输出通道：写 stdout 用的，线程安全
    stdout: Arc<Mutex<std::io::Stdout>>,
    /// 接收器（从 stdin 中读到的消息，通过广播发出）
    receiver: broadcast::Receiver<McpResult<McpMessage>>,
}

impl McpStdioTransport {
    pub fn new() -> (Self, broadcast::Sender<McpResult<McpMessage>>) {
        // broadcast 是 Tokio 提供的 多订阅者广播通道，可以有多个接受者（通过 clone 复制）
        let (sender, receiver) = broadcast::channel(100);
        let transport = Self {
            stdout: Arc::new(Mutex::new(std::io::stdout())),
            receiver,
        };

        // 启动后台任务：读取 stdin 并转成 MCP 消息广播出去
        let stdin = tokio::io::stdin();
        let mut reader = TokioBufReader::new(stdin);
        let sender_clone = sender.clone();

        // 就像多线程的 spawn 一样，注册了这个过程，一直在运作，直到这个 loop 被打破
        tokio::spawn(async move {
            let mut line = String::new();
            loop {
                line.clear();
                // 等待读取成功
                match reader.read_line(&mut line).await {
                    Ok(0) => break,
                    Ok(_) => {
                        // 得到了一行输入，将字符串视为 json 格式，并解释为 McpMessage
                        let message = match serde_json::from_str(&line) {
                            Ok(message) => Ok(message),
                            Err(err) => Err(McpError::Serialization(err.to_string())),
                        };
                        
                        // 使用 多订阅者广播通道 中的发送者发送这个 Message
                        // 那么其他 receiver 就应该可以接收到！
                        if sender_clone.send(message).is_err() {
                            break;
                        }
                    }
                    Err(err) => {
                        let _ = sender_clone.send(Err(McpError::Io(err.to_string())));
                    }
                }
            }
        });

        (transport, sender)
    }
}

#[async_trait]
impl McpTransport for McpStdioTransport {
    async fn send(&self, message: McpMessage) -> McpResult<()> {
        // 尝试获取 stdout 的控制权
        let mut stdout = self.stdout.lock().map_err(|_e| {
            McpError::protocol(McpErrorCode::InternalError, "Failed to acquire stdout lock")
        })?;

        // 将 message 解析为 json 字符串
        let json = serde_json::to_string(&message)?;

        // 调用 stdout 的 write 方法写入标准输出，并立刻刷新，将 stdout 重定向到了一 Mcp 端的输入就可以通信
        writeln!(stdout, "{}", json).map_err(|e| McpError::Io(e.to_string()))?;
        stdout.flush().map_err(|e| McpError::Io(e.to_string()))?;

        Ok(())
    }
    
    fn receive(&self) -> Pin<Box<dyn Stream<Item = McpResult<McpMessage>> + Send>> {
        // 每次被调用时，都会返回一个新的、可以异步消费的消息流。
        // 流中的每一项就是一条从 stdin 接收到的、被解析成 McpMessage 的消息。
        // 如果出错（比如通道关闭），这个流会结束。
        let rx = self.receiver.resubscribe();
        Box::pin(futures::stream::unfold(rx, |mut rx| async move {
            match rx.recv().await {
                Ok(msg) => Some((msg, rx)),
                Err(_) => None,
            }
        }))
    }
    
    async fn close(&self) -> McpResult<()> {
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::protocol::{McpRequest, McpRequestId};
    use std::sync::mpsc;
    use std::thread;
    use std::time::Duration;
    use tokio::sync::broadcast;
    use tokio::runtime::Runtime;

    #[test] 
    fn test_stdio_transport() {
        // Create a oneshot channel to verify message sending
        let (verify_tx, verify_rx) = mpsc::channel();
        
        // Spawn a thread to handle the transport operations
        thread::spawn(move || {
            // Create a tokio runtime for async operations
            let rt = Runtime::new().unwrap();
            
            // Create test message
            let request: McpRequest = McpRequest::new(
                "test_method",
                Some(serde_json::json!({"key": "value"})),
                McpRequestId::Number(1),
            );
            let message = McpMessage::Request(request);

            // Create transport with broadcast channel
            let (_, receiver) = broadcast::channel(100);
            let transport = McpStdioTransport {
                stdout: Arc::new(Mutex::new(std::io::stdout())),
                receiver,
            };

            // Test sending a message
            let send_result = rt.block_on(transport.send(message.clone()));
            verify_tx.send(send_result.is_ok()).unwrap();
        });

        // Wait for the result with timeout
        match verify_rx.recv_timeout(Duration::from_secs(1)) {
            Ok(true) => (), // Test passed
            Ok(false) => panic!("Failed to send message"),
            Err(_) => panic!("Test timed out"),
        }
    }
}